from db_init import get_db, create_tables
from models import User, ImageUpload, Task
from datetime import datetime
from sqlalchemy.exc import IntegrityError

def main():
    # 确保数据库表已创建
    create_tables()
    
    # 获取数据库会话
    db = next(get_db())
    
    try:
        # 创建一个用户
        print("Creating a test user...")
        try:
            # 尝试创建用户，如果已存在则跳过
            user = User(
                username="testuser",
                email="test@example.com",
                password_hash="hashed_password_here",  # 在实际应用中应使用密码哈希
                is_active=True
            )
            db.add(user)
            db.commit()
            db.refresh(user)
            print(f"User created: {user.username}")
        except IntegrityError:
            db.rollback()
            user = db.query(User).filter(User.username == "testuser").first()
            print(f"User already exists: {user.username}")
        
        # 上传图片
        print("\nUploading a test image...")
        image = ImageUpload(
            user_id=user.id,
            filename="example.jpg",
            filepath="/uploads/example.jpg",
            file_size=1024 * 512,  # 512KB
            mime_type="image/jpeg",
            description="Test image for demonstration",
            is_processed=False
        )
        db.add(image)
        db.commit()
        db.refresh(image)
        print(f"Image uploaded: {image.filename}")
        
        # 创建任务
        print("\nCreating a task...")
        task = Task(
            user_id=user.id,
            image_upload_id=image.id,
            task_type="classification",
            status="pending",
            result="",
            error_message=None
        )
        db.add(task)
        db.commit()
        db.refresh(task)
        print(f"Task created: {task.task_type} (Status: {task.status})")
        
        # 更新任务状态
        print("\nUpdating task status...")
        task.status = "completed"
        task.result = '{"classification": "cat", "confidence": 0.95}'
        task.completed_at = datetime.utcnow()
        db.commit()
        print(f"Task updated: {task.task_type} (Status: {task.status})")
        
        # 查询演示
        print("\nQuery demonstration:")
        
        # 获取用户的所有图片
        user_images = db.query(ImageUpload).filter(ImageUpload.user_id == user.id).all()
        print(f"User has {len(user_images)} images")
        
        # 获取图片的所有任务
        image_tasks = db.query(Task).filter(Task.image_upload_id == image.id).all()
        print(f"Image has {len(image_tasks)} tasks")
        
        # 获取已完成的任务
        completed_tasks = db.query(Task).filter(Task.status == "completed").all()
        print(f"There are {len(completed_tasks)} completed tasks")
        
    except Exception as e:
        print(f"An error occurred: {e}")
        db.rollback()
    finally:
        db.close()
        print("\nDatabase connection closed.")

if __name__ == "__main__":
    main()